# -*- coding: utf-8 -*-
"""
Created on Fri Aug 25 13:46:31 2017
VOC to COCO
@author: zxl
"""


import os
import sys
import shutil
import argparse
import json
import xml.etree.ElementTree as ET


def coco_to_voc(voc_root_path, coco_save_path, labels):
    """

    :param voc_root_path:
    :param coco_save_path:
    :param labels:
    :return:
    """
    coco_annotation_dir = os.path.join(coco_save_path, "annotations")
    if not os.path.exists(coco_annotation_dir):
        os.makedirs(coco_annotation_dir)

    coco_img_dir = os.path.join(coco_save_path, "train2017")
    if not os.path.exists(coco_img_dir):
        os.mkdir(coco_img_dir)

    json_file_train = os.path.join(coco_annotation_dir, 'instances_train.json')
    json_file_train_open_w = open(json_file_train, 'w')
    json_file_val = os.path.join(coco_annotation_dir, 'instances_val.json')
    json_file_val_open_w = open(json_file_val, 'w')  # TODO, CLOSE
    img_id = 0
    annotation_id = 0

    train_text = dict()
    train_text["info"] = {"description": "This is a train dataset "
                                         "of vehicle from DETRAC_PASCAL_v2 "
                                         "and CROWDAI_PASCAL set",
                          "url": "http://mscoco.org",
                          "version": "1.0",
                          "year": 2017,
                          "contributor": "Microsoft COCO group",
                          "date_created": "Fri Aug 25 13:46:31 2017"}

    train_text["licenses"] = [{"url": "http://detrac-db.rit.albany.edu/",
                               "name": "Attribution-NonCommercial-ShareAlike License",
                               "id": 1}]
    train_text["images"] = []
    train_text["annotations"] = []
    train_text["categories"] = []
    for i in range(len(labels)):
        obj_id = i + 1
        category = {"name": labels[i],
                    "supercategory": labels[i],
                    "id": obj_id}
        train_text["categories"].append(category)

    voc_train_filelist = open(os.path.join(voc_root_path,
                                           'ImageSets/Main/train.txt')).readlines()
    for line in voc_train_filelist:
        img_id += 1
        if img_id % 100 == 0:
            print("processing the {}th image in train set".format(img_id))
        xml_path = os.path.join(voc_root_path, 'Annotations/' + line.strip() + '.xml')
        xml_file = open(xml_path, "r")
        tree = ET.parse(xml_file)
        xml_file.close()
        root = tree.getroot()
        img_name = root.find("filename").text
        scr_img_file_full_path = os.path.join(voc_root_path, 'JPEGImages/' + img_name)
        dest_img_file_full_path = os.path.join(coco_img_dir, "" + img_name)
        if not os.path.exists(dest_img_file_full_path):
            shutil.copyfile(scr_img_file_full_path, dest_img_file_full_path)

        size = root.find('size')
        img_w = int(size.find('width').text)
        img_h = int(size.find('height').text)
        img = {"coco_url": "coco/train2017/" + str(img_name),
               "flickr_url": "coco/train2017/" + str(img_name),
               "file_name": img_name,
               "width": img_w,
               "height": img_h,
               "date_captured": "Fri Aug 25 13:46:31 2017",
               "license": 1,
               "id": img_id}
        train_text["images"].append(img)
        for obj in root.iter('object'):
            annotation_id += 1
            obj_class = obj.find('name').text
            obj_class_id = labels.index(obj_class) + 1  # class id start from 1

            xml_box = obj.find('bndbox')
            left = float(xml_box.find('xmin').text)
            top = float(xml_box.find('ymin').text)
            right = float(xml_box.find('xmax').text)
            bottom = float(xml_box.find('ymax').text)
            bbox_w = right - left
            bbox_h = bottom - top
            area = bbox_w * bbox_h
            annotation_coco = {"id": annotation_id,
                               "image_id": img_id,
                               "category_id": obj_class_id,
                               "segmentation": [[left, top,
                                                 left, bottom,
                                                 right, bottom,
                                                 right, top]],
                               "area": area,
                               "bbox": [left, top, bbox_w, bbox_h],
                               "iscrowd": 0}
            train_text["annotations"].append(annotation_coco)

    json_file_train_open_w.write(json.dumps(train_text, sort_keys=True,
                                            indent=4, separators=(', ', ': ')))
    json_file_train_open_w.close()
    print("%d train images have been completed" % img_id)
    train_image_num = img_id

    validate_text = dict()
    validate_text["info"] = {"description": "This is a validate"
                                            " dataset of vehicle just from DETRAC_PASCAL_v2",
                             "url": "http://mscoco.org", "version": "1.0",
                             "year": 2017,
                             "contributor": "Microsoft COCO group",
                             "date_created": "Fri Aug 25 13:46:31 2017"}

    validate_text["licenses"] = [{"url": "http://detrac-db.rit.albany.edu/",
                                  "name": "Attribution-NonCommercial-ShareAlike License",
                                  "id": 1}]
    validate_text["images"] = []
    validate_text["annotations"] = []
    validate_text["categories"] = []
    for i in range(len(labels)):
        obj_id = i + 1
        category = {"name": labels[i],
                    "supercategory": labels[i],
                    "id": obj_id}
        validate_text["categories"].append(category)

    voc_validate_filelist = open(os.path.join(voc_root_path, 'ImageSets/Main/test.txt')).readlines()
    for line in voc_validate_filelist:
        img_id += 1
        if (img_id - train_image_num) % 100 == 0:
            print("processing the {}th image in validate set".format(int(img_id-train_image_num)))
        xml_path = os.path.join(voc_root_path, 'Annotations/'+line.strip()+'.xml')
        xml_file = open(xml_path, 'r')
        tree = ET.parse(xml_file)
        xml_file.close()
        root = tree.getroot()
        img_name = root.find("filename").text

        scr_file = os.path.join(voc_root_path, 'JPEGImages/' + img_name)
        target_file = os.path.join(coco_save_path, 'val2017/' + img_name)
        if not os.path.exists(os.path.join(coco_save_path, "val2017")):
            os.makedirs(os.path.join(coco_save_path, "val2017"))
        if not os.path.exists(target_file):
            shutil.copyfile(scr_file, target_file)

        size = root.find('size')
        img_w = int(size.find('width').text)
        img_h = int(size.find('height').text)
        img = {"coco_url": "coco/validate2017/"+str(img_name),
               "flickr_url": "coco/validate2017/" + str(img_name),
               "file_name": img_name,
               "width": img_w,
               "height": img_h,
               "date_captured": "Fri Aug 25 13:46:31 2017",
               "license": 1,
               "id": img_id}
        validate_text["images"].append(img)
        for obj in root.iter('object'):
            annotation_id += 1
            obj_class = obj.find('name').text
            obj_class_id = labels.index(obj_class)  # class id start from 0

            xml_box = obj.find('bndbox')
            left = float(xml_box.find('xmin').text)
            top = float(xml_box.find('ymin').text)
            right = float(xml_box.find('xmax').text)
            bottom = float(xml_box.find('ymax').text)
            bbox_w = right - left
            bbox_h = bottom - top
            area = bbox_w * bbox_h
            annotation_coco = {"id": annotation_id,
                               "image_id": img_id,
                               "category_id": obj_class_id,
                               "segmentation": [[left, top,
                                                 left, bottom,
                                                 right, bottom,
                                                 right, top]],
                               "area": area,
                               "bbox": [left, top, bbox_w, bbox_h],
                               "iscrowd": 0}
            validate_text["annotations"].append(annotation_coco)

    json_file_val_open_w.write(json.dumps(validate_text, sort_keys=True,
                                          indent=4, separators=(', ', ': ')))
    json_file_val_open_w.close()
    print("{} validate images have been completed".format(int(img_id-train_image_num)))


if __name__ == '__main__':
    try:
        VOC_ROOT_DIR = sys.argv[1]
        COCO_SAVE_DIR = sys.argv[2]
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--voc_dir", type=str,
                            default=os.path.join(os.getenv("HOME"),
                                                 "PycharmProjects/data_tool/data/voc2012"),
                            help="VOC directory")
        PARSER.add_argument("--coco_dir", type=str,
                            default=os.path.join(os.getenv("HOME"),
                                                 "PycharmProjects/data_tool/data/coco"),
                            help="COCO directory")

        ARGS = PARSER.parse_args()
        VOC_ROOT_DIR = ARGS.voc_dir
        COCO_SAVE_DIR = ARGS.coco_dir
    LABELS = ['person', 'bicycle', 'car', 'motorbike',
              'aeroplane', 'bus', 'train', 'truck', 'boat',
              'traffic light', 'fire hydrant', 'stop sign',
              'parking meter', 'bench', 'bird', 'cat', 'dog',
              'horse', 'sheep', 'cow', 'elephant', 'bear',
              'zebra', 'giraffe', 'backpack', 'umbrella',
              'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
              'snowboard', 'sports ball', 'kite', 'baseball bat',
              'baseball glove', 'skateboard', 'surfboard',
              'tennis racket', 'bottle', 'wine glass', 'cup',
              'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
              'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog',
              'pizza', 'donut', 'cake', 'chair', 'sofa', 'pottedplant',
              'bed', 'diningtable', 'toilet', 'tvmonitor', 'laptop',
              'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
              'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
              'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

    coco_to_voc(VOC_ROOT_DIR, COCO_SAVE_DIR, LABELS)
